import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


# img_dir = '../../../../../large_data/CV2/lesson/Day07'
# img_name = 'football_ball.jpg'
# temp_path = os.path.join(img_dir, img_name)

img_dir = '../../../../large_data/pic/'
# img_name = 'DSC05039.JPG'
# img_name = 'DSC05022_1.JPG'
img_name = 'dog.jpg'
# img_name = 'dog_bird.jpg'
# img_name = 'football.jpg'
# img_name = 'messi5.jpg'
# img_dir = '../../../../large_data/pic/watershed/'
# img_name = 'water_coins.jpg'
img_path = os.path.join(img_dir, img_name)

spr = 2
spc = 2
spn = 0
plt.figure(figsize=[14, 7])

sep('load')
img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)
print('original shape', img.shape)
H, W = img.shape
H2 = H // 2
W2 = W // 2
print(W, H)
my_show_img(img, 'original by cv', cmap='gray')

sep('fourier')
img = img.astype(np.float32)
f = cv.dft(img, flags=cv.DFT_COMPLEX_OUTPUT)
print_numpy_ndarray_info(f, 'f')
fshift = np.fft.fftshift(f)
print_numpy_ndarray_info(fshift, 'fshift')
fshift_ = fshift.copy()

sep('filter HPF')
SIZE = 30
fshift[H2 - SIZE:H2 + SIZE, W2 - SIZE: W2 + SIZE] = 0

sep('inverse')
fi = np.fft.ifftshift(fshift)
imgi = cv.idft(fi)
imgi = cv.magnitude(imgi[:, :, 0], imgi[:, :, 1])
my_show_img(imgi, 'HPF', cmap='gray')

sep('filter LPF')
mask = np.zeros((H, W, 2), dtype=np.float32)
mask[H2 - SIZE:H2 + SIZE, W2 - SIZE: W2 + SIZE] = 1.
fshift = fshift_.copy()
fshift *= mask

sep('inverse')
fi = np.fft.ifftshift(fshift)
imgi = cv.idft(fi)
imgi = cv.magnitude(imgi[:, :, 0], imgi[:, :, 1])
my_show_img(imgi, 'LPF', cmap='gray')

sep('show')
plt.show()
